simion.workbench_program()

local kineticEnergyTable = {1}
local L1_voltageTable    = {1}
local L2_voltageTable    = {1}
local L3_voltageTable    = {1}
local L4_voltageTable    = {1}
local KineticEnergyIndex = 1
local NUMENGs = 1
local readin = true

--read the initial lens table
local voltage_file = assert(io.open("Voltages.dat","r"))
while readin do
  kineticEnergyTable[NUMENGs] = voltage_file:read("*n")
  L1_voltageTable[NUMENGs] = voltage_file:read("*n")
  L2_voltageTable[NUMENGs] = voltage_file:read("*n")
  L3_voltageTable[NUMENGs] = voltage_file:read("*n")
  L4_voltageTable[NUMENGs] = voltage_file:read("*n")
  readin = (kineticEnergyTable[NUMENGs] ~= nil)
  NUMENGs = NUMENGs + 1
end
NUMENGs = #kineticEnergyTable
print(NUMENGs)
voltage_file:close()

KineticEnergyIndex = 1
--local KineticEnergy = kineticEnergyTable[KineticEnergyIndex]
local KineticEnergy
local maxkineticEnergy = 7000
local PassEnergy    = 250
local EnergyStep = 5
local StepControl

adjustable _L1_voltage
adjustable _L2_voltage
adjustable _L3_voltage
adjustable _L4_voltage
local L1_step
local L2_step
local L3_step
local L4_step
adjustable _Ad_voltage
local _Ae_voltage
local _Ai_voltage
local Ad_step
adjustable _opt_minradius

local firstrun = 1

local SimplexOptimizer = require "simionx.SimplexOptimizer"
local opt
local NUM_maxcalls = 1000

local total_ion_number = 0
local splat_ion_number = 0
--The numbering of the lens electrodes
local L01 = 5
local L02 = 6
local L1e = 7
local L1i = 8
local L2e = 9
local L2i = 10
local L3e = 11
local L3i = 12
local L4e = 13
local L4i = 13
local L5  = 14
local Ae  = 15
local Ai  = 16
local D1e = 2
local D1i = 4
local D2e = 1
local D2i = 3
local Det = 17

local run_number = 0
local is_optimizing = 1
local startopt = 1

local opt_output_file = assert(io.open("scan.dat","w"))
opt_output_file:write(string.format("run, KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, _Ad_voltage, transmission, y_error, metric\n"))
local final_ouput_file = assert(io.open("final.dat","w"))
--final_ouput_file:write(string.format("ion_number, tof_us, y_mm\n"))
--final_ouput_file:write(string.format("KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, _Ae_voltage, _Ai_voltage, transmission, y_error\n"))
final_ouput_file:write(string.format("KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, transmission, y_error\n"))
local initial_ouput_file = assert(io.open("initial.dat","w"))
--initial_ouput_file:write(string.format("KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, _Ae_voltage, _Ai_voltage, transmission, y_error\n"))
initial_ouput_file:write(string.format("KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, transmission, y_error\n"))
local lens_250_ouput_file = assert(io.open("lens250.txt","w"))
lens_250_ouput_file:write(string.format("[Data] \n"))
lens_250_ouput_file:write(string.format("KE   L1e   L1i   L2e   L2i   L3e   L3i   L4e   L4i \n"))
local lens_120_ouput_file = assert(io.open("lens120.txt","w"))
lens_120_ouput_file:write(string.format("[Data] \n")) 
lens_120_ouput_file:write(string.format("KE   L1e   L1i   L2e   L2i   L3e   L3i   L4e   L4i \n"))
local lens_80_ouput_file = assert(io.open("lens80.txt","w"))
lens_80_ouput_file:write(string.format("[Data] \n")) 
lens_80_ouput_file:write(string.format("KE   L1e   L1i   L2e   L2i   L3e   L3i   L4e   L4i \n"))
local lens_60_ouput_file = assert(io.open("lens60.txt","w"))
lens_60_ouput_file:write(string.format("[Data] \n")) 
lens_60_ouput_file:write(string.format("KE   L1e   L1i   L2e   L2i   L3e   L3i   L4e   L4i \n"))
local lens_40_ouput_file = assert(io.open("lens40.txt","w"))
lens_40_ouput_file:write(string.format("[Data] \n")) 
lens_40_ouput_file:write(string.format("KE   L1e   L1i   L2e   L2i   L3e   L3i   L4e   L4i \n"))
local lens_20_ouput_file = assert(io.open("lens20.txt","w"))
lens_20_ouput_file:write(string.format("[Data] \n")) 
lens_20_ouput_file:write(string.format("KE   L1e   L1i   L2e   L2i   L3e   L3i   L4e   L4i \n"))
local lens_10_ouput_file = assert(io.open("lens10.txt","w"))
lens_10_ouput_file:write(string.format("[Data] \n")) 
lens_10_ouput_file:write(string.format("KE   L1e   L1i   L2e   L2i   L3e   L3i   L4e   L4i \n"))
local lens_5_ouput_file = assert(io.open("lens5.txt","w"))
lens_5_ouput_file:write(string.format("[Data] \n")) 
lens_5_ouput_file:write(string.format("KE   L1e   L1i   L2e   L2i   L3e   L3i   L4e   L4i \n"))

local max_y = -10000
local min_y =  10000

local sum_y = 0
local sum_y2 = 0
local averaged_y = 0
local error_y = 0
local transmission = 0
local metric = 0
local newL1, newL2, newL3, newL4, newtrans, newerr, newmetric

function init_lens_voltage()
	--_L1_voltage =  -26 + 37.068 + 0.88529 *KineticEnergy + 0.00020981 *KineticEnergy*KineticEnergy - 2.7546e-008 *KineticEnergy*KineticEnergy*KineticEnergy
	--_L2_voltage =  -33 + 50.833 + 2.0386 *KineticEnergy + 0.00023757*KineticEnergy*KineticEnergy - 3.1899e-008 *KineticEnergy*KineticEnergy*KineticEnergy
	--_L3_voltage =   2  + 154.18 + 0.25616*KineticEnergy + 2.2802e-005 *KineticEnergy*KineticEnergy + 3.7177e-010 *KineticEnergy*KineticEnergy*KineticEnergy
	--_L4_voltage =  -58 + 384.34 + 0.27737*KineticEnergy - 8.3885e-005 *KineticEnergy*KineticEnergy + 5.8825e-009 *KineticEnergy*KineticEnergy*KineticEnergy
	KineticEnergy = kineticEnergyTable[KineticEnergyIndex]
	_L1_voltage =  L1_voltageTable[KineticEnergyIndex]
  _L2_voltage =  L2_voltageTable[KineticEnergyIndex]
	_L3_voltage =  L3_voltageTable[KineticEnergyIndex]
	_L4_voltage =  L4_voltageTable[KineticEnergyIndex]
	if((run_number ~= 0) and (firstrun == 1)) then
		_L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage = newL1, newL2, newL3, newL4
	end
	
	--_L1_voltage = math.max(_L1_voltage,1.25*KineticEnergy)
	--_L3_voltage = math.min(_L3_voltage,2.5*math.sqrt(KineticEnergy*PassEnergy))
	--if _L3_voltage<_L4_voltage then
	--  _L3_voltage, _L4_voltage = _L4_voltage, _L3_voltage
	--end
	
  StepControl = KineticEnergy/2
  _opt_minradius = math.min(KineticEnergy*0.004,0.5)
  L1_step = _L1_voltage*0.05
  L2_step = _L2_voltage*0.05
  L3_step = _L3_voltage*0.05
  L4_step = _L4_voltage*0.02
  _Ad_voltage = 10.9375
  _Ae_voltage = PassEnergy - _Ad_voltage
  _Ai_voltage = PassEnergy + _Ad_voltage
  Ad_step = PassEnergy/100
end

init_lens_voltage()

function segment.initialize()
  if (run_number ==0) then
    sim_trajectory_quality = 3
    --sim_trajectory_image_control = 3
    total_ion_number = ion_number -- get the total_ion_number
  end
  -- Change the kinetic of electron
  local ionspeed, ionaz, ionel = rect3d_to_polar3d(ion_vx_mm, ion_vy_mm, ion_vz_mm)
  --print(ion_vx_mm, ion_vy_mm, ion_vz_mm)
  --local old_ion_vx_mm = ion_vx_mm
  ionspeed = ke_to_speed(KineticEnergy, ion_mass)
  ion_vx_mm, ion_vy_mm, ion_vz_mm = polar3d_to_rect3d(ionspeed, ionaz, ionel)
  --print(ion_vx_mm, ion_vy_mm, ion_vz_mm)
  --print((ion_vx_mm/old_ion_vx_mm)^2)
  sim_rerun_flym = is_optimizing
  --sim_trajectory_image_control = 0
end

function segment.init_p_values()
  if (startopt == 1) then
  	--if (KineticEnergyIndex <= NUMENGs) then
    	--init_lens_voltage()
  	--end
    --opt = SimplexOptimizer {
    --  start = {_L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, _Ad_voltage}, 
    --  step = {L1_step, L2_step, L3_step, L4_step, Ad_step},
    --  maxcalls = NUM_maxcalls, minradius = _opt_minradius
    --}
    opt = SimplexOptimizer {
      start = {_L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage}, 
      step = {L1_step, L2_step, L3_step, L4_step},
      maxcalls = NUM_maxcalls, minradius = _opt_minradius
    }
    startopt = 0
  end
  --_L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, _Ad_voltage = opt:values() 
  _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage = opt:values()
  --print(_L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage)
  --_Ae_voltage = PassEnergy - _Ad_voltage
  --_Ai_voltage = PassEnergy + _Ad_voltage
  adj_elect[L01] = KineticEnergy
  adj_elect[L02] = KineticEnergy
  adj_elect[L1e] = _L1_voltage
  adj_elect[L1i] = _L1_voltage
  adj_elect[L2e] = _L2_voltage
  adj_elect[L2i] = _L2_voltage
  adj_elect[L3e] = _L3_voltage
  adj_elect[L3i] = _L3_voltage
  adj_elect[L4e] = _L4_voltage
  adj_elect[L4i] = _L4_voltage
  adj_elect[L5]  = PassEnergy
  adj_elect[Ae]  = _Ae_voltage
  adj_elect[Ai]  = _Ai_voltage
  adj_elect[D1e] = 201.75
	adj_elect[D1i] = 303.25
	adj_elect[D2e] = 303.25
	adj_elect[D2i] = 201.75
	adj_elect[Det] = PassEnergy
  --print("init_p_values")
end

function segment.terminate()
  --sim_retain_changed_potentials = 1
  if (sim_rerun_flym == 1) then --record the optimized condition
    --print(string.format("ion_px_mm = %f", ion_px_mm))
    --if (true) then
    if (ion_px_mm < 5) then
      splat_ion_number = splat_ion_number + 1
      sum_y = sum_y + ion_py_mm
      sum_y2 = sum_y2 + ion_py_mm * ion_py_mm
      max_y = max(max_y, ion_py_mm)
      min_y = min(min_y, ion_py_mm)
      --print(ion_number, ion_py_mm, sum_y, sum_y2)
    end
    
    if ( ion_number == total_ion_number) then
      averaged_y = sum_y / splat_ion_number
      error_y = math.sqrt((sum_y2 - splat_ion_number * averaged_y * averaged_y)/splat_ion_number)
      --error_y = max_y - min_y
      transmission = splat_ion_number/total_ion_number
      --print(string.format("splat_ion_number = %d, total_ion_number = %d", splat_ion_number, total_ion_number))
      run_number = opt:ncalls()
      metric = 1/transmission/transmission*error_y
      opt:result(metric)
      opt_output_file:write(string.format("%d, %f, %f, %f, %f, %f, %f, %f, %f, %f \n", run_number, KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, _Ad_voltage, transmission, error_y, metric))
      if (run_number == NUM_maxcalls) then
        opt_output_file:write(string.format("out of iteration limit: %d \n", NUM_maxcalls))
      end
      opt_output_file:flush()
            
      sum_y = 0
      sum_y2 = 0
      max_y = -10000
      min_y =  10000
      splat_ion_number = 0
      
      if (run_number == 1) then
        --initial_ouput_file:write(string.format("%d, %f, %f, %f, %f, %f, %f, %f, %f \n", KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, _Ae_voltage, _Ai_voltage, transmission, error_y))
        initial_ouput_file:write(string.format("%d, %f, %f, %f, %f, %f, %f \n", KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, transmission, error_y))
        initial_ouput_file:flush()
        newL1 = _L1_voltage
        newL2 = _L2_voltage
        newL3 = _L3_voltage
        newL4 = _L4_voltage
        newtrans = transmission
        newerr = error_y
        newmetric = metric
      elseif (newmetric > metric and metric > 1e-9) then
        newL1 = _L1_voltage
        newL2 = _L2_voltage
        newL3 = _L3_voltage
        newL4 = _L4_voltage
        newtrans = transmission
        newerr = error_y
        newmetric = metric
      end
      is_optimizing = opt:running() and 1 or 0
      if (is_optimizing == 0 or metric < 1e-9 or error_y < 1e-9 or transmission < 1e-9) then
        --final_ouput_file:write(string.format("%d, %f, %f, %f, %f, %f, %f, %f, %f \n", KineticEnergy, _L1_voltage, _L2_voltage, _L3_voltage, _L4_voltage, _Ae_voltage, _Ai_voltage, transmission, error_y))
        final_ouput_file:write(string.format("%f, %f, %f, %f, %f, %f, %f \n", KineticEnergy, newL1, newL2, newL3, newL4, newtrans, newerr))
        if(metric < 1e-9 or error_y < 1e-9 or transmission < 1e-9) then
          final_ouput_file:write(string.format("metric or error_y or transmission is unphysically zero \n"))
        else
          lens_250_ouput_file:write(string.format("%f   %f   %f   %f   %f   %f   %f   %f   %f \n", KineticEnergy,
            newL1-KineticEnergy, newL1-KineticEnergy, newL2-KineticEnergy, newL2-KineticEnergy,
            newL3-KineticEnergy, newL3-KineticEnergy, newL4-KineticEnergy, newL4-KineticEnergy ))
          lens_250_ouput_file:flush()
          lens_120_ouput_file:write(string.format("%f   %f   %f   %f   %f   %f   %f   %f   %f \n", KineticEnergy/25*12,
            ( newL1-KineticEnergy)/25*12, ( newL1-KineticEnergy)/25*12, ( newL2-KineticEnergy)/25*12, ( newL2-KineticEnergy)/25*12,
            ( newL3-KineticEnergy)/25*12, ( newL3-KineticEnergy)/25*12, ( newL4-KineticEnergy)/25*12, ( newL4-KineticEnergy)/25*12 ))
          lens_120_ouput_file:flush()
          lens_80_ouput_file:write(string.format("%f   %f   %f   %f   %f   %f   %f   %f   %f \n", KineticEnergy/25*8,
            ( newL1-KineticEnergy)/25*8, ( newL1-KineticEnergy)/25*8, ( newL2-KineticEnergy)/25*8, ( newL2-KineticEnergy)/25*8,
            ( newL3-KineticEnergy)/25*8, ( newL3-KineticEnergy)/25*8, ( newL4-KineticEnergy)/25*8, ( newL4-KineticEnergy)/25*8 ))
          lens_80_ouput_file:flush()
          lens_60_ouput_file:write(string.format("%f   %f   %f   %f   %f   %f   %f   %f   %f \n", KineticEnergy/25*6,
            ( newL1-KineticEnergy)/25*6, ( newL1-KineticEnergy)/25*6, ( newL2-KineticEnergy)/25*6, ( newL2-KineticEnergy)/25*6,
            ( newL3-KineticEnergy)/25*6, ( newL3-KineticEnergy)/25*6, ( newL4-KineticEnergy)/25*6, ( newL4-KineticEnergy)/25*6 ))
          lens_60_ouput_file:flush()  
          lens_40_ouput_file:write(string.format("%f   %f   %f   %f   %f   %f   %f   %f   %f \n", KineticEnergy/25*4,
            ( newL1-KineticEnergy)/25*4, ( newL1-KineticEnergy)/25*4, ( newL2-KineticEnergy)/25*4, ( newL2-KineticEnergy)/25*4,
            ( newL3-KineticEnergy)/25*4, ( newL3-KineticEnergy)/25*4, ( newL4-KineticEnergy)/25*4, ( newL4-KineticEnergy)/25*4 ))
          lens_40_ouput_file:flush()
          lens_20_ouput_file:write(string.format("%f   %f   %f   %f   %f   %f   %f   %f   %f \n", KineticEnergy/25*2,
            ( newL1-KineticEnergy)/25*2, ( newL1-KineticEnergy)/25*2, ( newL2-KineticEnergy)/25*2, ( newL2-KineticEnergy)/25*2,
            ( newL3-KineticEnergy)/25*2, ( newL3-KineticEnergy)/25*2, ( newL4-KineticEnergy)/25*2, ( newL4-KineticEnergy)/25*2 ))
          lens_20_ouput_file:flush()
          lens_10_ouput_file:write(string.format("%f   %f   %f   %f   %f   %f   %f   %f   %f \n", KineticEnergy/25*1,
            ( newL1-KineticEnergy)/25*1, ( newL1-KineticEnergy)/25*1, ( newL2-KineticEnergy)/25*1, ( newL2-KineticEnergy)/25*1,
            ( newL3-KineticEnergy)/25*1, ( newL3-KineticEnergy)/25*1, ( newL4-KineticEnergy)/25*1, ( newL4-KineticEnergy)/25*1 ))
          lens_10_ouput_file:flush()
          lens_5_ouput_file:write(string.format("%f   %f   %f   %f   %f   %f   %f   %f   %f \n", KineticEnergy/25*0.5,
            ( newL1-KineticEnergy)/25*0.5, ( newL1-KineticEnergy)/25*0.5, ( newL2-KineticEnergy)/25*0.5, ( newL2-KineticEnergy)/25*0.5,
            ( newL3-KineticEnergy)/25*0.5, ( newL3-KineticEnergy)/25*0.5, ( newL4-KineticEnergy)/25*0.5, ( newL4-KineticEnergy)/25*0.5 ))
          lens_5_ouput_file:flush()
        end
        final_ouput_file:flush()
        --EnergyStep = ceil(EnergyStep * 1.2)
        --KineticEnergy = KineticEnergy + EnergyStep
        KineticEnergyIndex = KineticEnergyIndex +1
        --KineticEnergy = kineticEnergyTable[KineticEnergyIndex]
        --if (KineticEnergy < maxkineticEnergy) then 
        if (KineticEnergyIndex<=NUMENGs) then 
          startopt = 1
          is_optimizing = 1 
          init_lens_voltage()
        else
          --KineticEnergy = KineticEnergy - EnergyStep
          --KineticEnergy = kineticEnergyTable[KineticEnergyIndex-1]
          is_optimizing = 0
        end
      end
    end
  else
    --final_ouput_file:write(string.format("%d, %f, %f\n", ion_number, ion_time_of_flight, ion_py_mm))
    --final_ouput_file:flush()
    if (ion_number == total_ion_number) then     
      opt_output_file:close()
      initial_ouput_file:close()
      final_ouput_file:close()
      lens_250_ouput_file:close()
      lens_120_ouput_file:close()
      lens_80_ouput_file:close()
      lens_60_ouput_file:close()
      lens_40_ouput_file:close()
      lens_20_ouput_file:close()
      lens_10_ouput_file:close()
      lens_5_ouput_file:close()
      sim_retain_changed_potentials = 1
    end
  end
end

--opt_output_file:flush()